# ----------------------------------------------------
# kitti dataset api
# ----------------------------------------------------

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import os.path as osp
from .imdb import IMDB
# import xml.etree.ElementTree as ET
import numpy as np
import scipy.sparse
import cv2
import pickle
import subprocess
import uuid
import PIL
# TODO here need import evaluation protocal
import re

class kitti(IMDB):
    #def __init__(self, dataset_name, image_set, allowed_classes, use_diff=False):
    #def __init__(self, image_set, root_path, devkit_path, result_path, mask_size=-1, binary_thresh=None):
    def __init__(self, dataset_name, allowed_classes, devkit_path, result_path, mask_size=-1, binary_thresh=None):
        # 
        # image_set means training / testing
        # TODO, need confirm the dataset's name
        allowed_classes     = ['car']
        # suppose the dataset_name is kitti-det_training
        dataset_name, image_set     = dataset_name.split('_')[0:2]

        name = dataset_name+'_'+image_set
        root_path   = osp.join(osp.abspath(osp.dirname(osp.dirname(osp.dirname(__file__)))), 'data')
        # namely the './data/kitti'
        self._devkit_path       = devkit_path
        print("kitti's root_path is {:s}".format(root_path))

        super(kitti, self).__init__(name, image_set, root_path, devkit_path, result_path)

        self._dataset_name      = dataset_name
        self._image_set         = image_set
        self._dataset_path      = self._get_default_dataset_path()
        self._data_path         = self._get_default_data_path()
        self._kitti_classes     = ('__background__', \
                           'car', 'cyclist', 'pedestrian', 'misc'\
                           'person_sitting', 'tram', 'truck', 'van')
        # TODO
        self._allowed_classes   = allowed_classes
        self._allowed_str       = '_'.join(self._allowed_classes)
        self._classes           = self._get_allowed_classes(self._allowed_classes, self._kitti_classes)
        self._class_to_ind      = dict(list(zip(self._classes, list(range(len(self._classes))))))

        self.classes            = self._classes
        self.num_classes        = len(self.classes)
        # confirm the image file's ext
        self.check_img_ext()
        self._load_image_set_index()
        self._roidb_handler     = self.gt_roidb
        ## need run this to generate the image_index
        self._gt_roidb_call_num = 0
        self._roidb_handler()
        self.image_set_index    = self._image_index 
        self.num_images         = len(self._image_index)
        self._salt              = str(uuid.uuid4())
        self._comp_id           = 'comp4'
        # kitti specific config options
        self.config = {'cleanup': True, 'use_salt':False, \
                       'rpn_file': None}
        if result_path is None:
            subdir              = self._get_comp_id() + '_det_' + self._image_set +'_'+ self._allowed_str
            result_path         = osp.join(osp.join(self._data_path, 'results_deformConv'), subdir)
            self._result_path   = result_path

    def _get_allowed_classes(self, allowed_classes, all_classes):
        """
        using bbox with class in allowed_classes
        """
        if allowed_classes[0] == 'all':
            return all_classes
        else:
            for cls in allowed_classes:
                assert cls in all_classes, "{} in allowed_classes must in all_classes".format(cls)
            if '__background__' not in allowed_classes:
                tmp             = ['__background__']
                [tmp.append(cls) for cls in allowed_classes]
                allowed_classes = tuple(tmp)
            return allowed_classes
    
    def image_path_at(self, i):
        """
        return the absolute path to image i in the image sequence
        """
        return self.image_path_from_index(self._image_index[i])
    def image_path_from_index(self, index):
        """
        construct an image path from the image's "index" identifier
        """
        image_path              = os.path.join(self._data_path, 'image_2', 
                                    index + '.' + self._image_ext)
        assert osp.exists(image_path), "Path does not exists: {}".format(image_path)
        return image_path
    def _get_default_dataset_path(self):
        """
        return the default path where kitti is expected to exist 
        """
        tmp     = osp.join(self._devkit_path, self._dataset_name)
        assert osp.exists(tmp), "{} must exist".format(tmp)
        return tmp
    def _get_default_data_path(self):
        """
        return the default path where kitti-data is expected to exist
        """
        tmp     = osp.join(self._dataset_path, self._image_set)
        assert osp.exists(tmp), "{} must exist".format(tmp)
        return tmp
    def check_img_ext(self):
        """
        check the image extension name for kitti, 'jpg' or 'png' 
        """
        img_dir     = osp.join(self._data_path, 'image_2')
        img_ext_list= ['jpg', 'png']
        for img_ext in img_ext_list:
            one_img_path    = os.path.join(img_dir, '000000.'+img_ext)
            if osp.exists(one_img_path):
                self._image_ext = img_ext
                return True
        assert False, "Cannot find jpg/png file {}".format(img_dir)
    def _load_image_set_index(self):
        """
        Load the indexes listed in this dataset's image set file
        """
        self._image_index_dir   = osp.join(self._data_path, 'Main')
        if not osp.exists(self._image_index_dir):
            os.makedirs(self._image_index_dir)
        self._image_index_file  = osp.join(self._image_index_dir, self._image_set + '_' + \
            self._allowed_str+ '.txt')
        self._image_index_size_file = osp.join(self._image_index_dir, self._image_set + '_' + self._allowed_str + '_size.pkl')
        if osp.exists(self._image_index_file) and osp.exists(self._image_index_size_file):
            # read image index from cache file
            with open(self._image_index_file, 'r') as f:
                _image_index    = [x.strip() for x in f.readlines()]
                # TODO
                # because this file is write bt self._confirm_image_index, so it is right
                self._image_index   = _image_index
            with open(self._image_index_size_file, 'rb') as f:
                _image_size_dict    = pickle.load(f)
                self._image_size_dict    = _image_size_dict
        else:
            tmp_list            = []
            size_dict           = {}
            _image_file_dir     = osp.join(self._data_path, 'image_2')
            for img_name in sorted(os.listdir(_image_file_dir)):
                pre_name, post_ext  = img_name.split('.')[0:2]
                if post_ext == self._image_ext:
                    tmp_list.append(pre_name)
                    img_size    = PIL.Image.open(self.image_path_from_index(pre_name)).size
                    img_size    = [img_size[0], img_size[1]]# width and height
                    size_dict[pre_name]     = img_size
#            _image_index        = tmp_list
            self._tmp_image_index   = tmp_list
            self.size_dict      = size_dict
            # write size cache to file
            with open(self._image_index_size_file, 'wb') as f:
                pickle.dump(size_dict, f, pickle.HIGHEST_PROTOCOL)
 #       return _image_index
        return None

    def gt_roidb(self):
        """
        Return the database of gt regions of interest
        """
        cache_file              = osp.join(self.cache_path, self.name+ '_'+ self._allowed_str+'_gt_roidb.pkl')
        if osp.exists(cache_file):
            with open(cache_file, 'rb') as fid:
                try:
                    roidb           = pickle.load(fid)
                except:
                    roidb           = pickle.load(fid, encoding='bytes')
            print('{} gt roidb loaded from {}'.format(self.name+ '_'+self._allowed_str, cache_file))
            return roidb
        else:
            if self._gt_roidb_call_num == 0:
                # load kitti gt info in the first time
                gt_roidb            = []
                _image_index        = []
                for index in self._tmp_image_index:
                    i_roidb         = self._load_kitti_annotation(index)
                    if i_roidb is not None:
                        gt_roidb.append(i_roidb)
                        _image_index.append(index)

                # confirm the _image_index
                self._confirm_image_index(_image_index)
                self._gt_roidb_call_num += 1
            else:
                # load kitti directly from self._image_index
                gt_roidb            = [self._load_kitti_annotation(index) for index in self._image_index]
            with open(cache_file, 'wb') as fid:
                pickle.dump(gt_roidb, fid, pickle.HIGHEST_PROTOCOL)
            print('wrote gt roidb to {}'.format(cache_file))
        return gt_roidb

    def _confirm_image_index(self, _image_index):
        """
        setting the attribute of class, and write into cache file
        """
        self._image_index           = _image_index
        with open(self._image_index_file, 'w') as fid:
            fid.writelines([str(ii)+'\n' for ii in self._image_index])

    def _load_kitti_annotation(self, index):
        """
        Load image and bounding boxes info from txt file in the kitti-det format
        """

        # use code from pascal_voc.load_pascal_annotation

        filename    = osp.join(self._data_path, 'label_2', index+'.txt')
        assert osp.exists(filename), '{} must exist'.format(filename)
        with open(filename, 'r') as f:
            anno_info_list = [x.strip() for x in f.readlines()]
            num_objs    = len(anno_info_list)
            assert num_objs > 0, "{} not consisting objects".format(filename)
            #boxes       = np.zeros((num_objs, 4), dtype=np.uint16)
            #gt_classes  = np.zeros((num_objs), dtype=np.uint32)
            #overlaps    = np.zeros((num_objs, self.num_classes), dtype=np.float32)

            box_list, gt_list, ol_list  = [], [], []
            for ix, anno_info in enumerate(anno_info_list):
                anno_info_sp    = re.split(r'\s+', anno_info)
                assert len(anno_info_sp) == 15, "kitti-detection must have 15 columns, but here is {}, length: {}".format(anno_info_sp, len(anno_info_sp))
                # Make pixel indexes 0-based, but kitti'gt already 0-based
                #x1      = float(anno_info_sp[4]) - 1
                #y1      = float(anno_info_sp[5]) - 1
                #x2      = float(anno_info_sp[6]) - 1
                #y2      = float(anno_info_sp[7]) - 1
                x1      = float(anno_info_sp[4])
                y1      = float(anno_info_sp[5])
                x2      = float(anno_info_sp[6])
                y2      = float(anno_info_sp[7])
                cls_name= str(anno_info_sp[0].lower())
                # 
                if cls_name in self._classes:
                    # if the cls_name in the allowed classes 
                    i_bbox  = np.array([[x1, y1, x2, y2]])
                    i_cls   = self._class_to_ind[cls_name]
                    i_overlap = np.zeros((1, self.num_classes), dtype=np.float32)
                    i_overlap[0, i_cls] = 1.0

                    box_list.append(i_bbox)
                    gt_list.append(i_cls)
                    ol_list.append(i_overlap)
            if len(box_list) > 0:
                # there exist allowed bbox
                boxes       = np.concatenate(box_list, 0)
                gt_classes  = np.array(gt_list)
                overlaps    = np.concatenate(ol_list, 0)

                # use code from pascal_voc.load_pascal_annotation
                roi_rec     = {}
                roi_rec['image']    = self.image_path_from_index(index)
#                size                = cv2.imread(roi_rec['image']).shape
                size                = self.size_dict[index]
                # need compatible with the PIL.Image
                roi_rec['height']   = size[1]
                roi_rec['width']    = size[0]
                roi_rec.update({'boxes': boxes, 
                        'gt_classes': gt_classes,
                        'gt_overlaps': overlaps,
                        'max_classes': overlaps.argmax(axis=1),
                        'max_overlaps':overlaps.max(axis=1),
                        'index': index, # my own attribute for debug
                        'flipped': False})
                return roi_rec
            else:
                return None

    def _load_rpn_roidb(self, gt_roidb):
        """
        """
        filename            = self.config['rpn_file']
        print('loading {}'.format(filename))
        assert os.path.exists(filename), \
            'rpn data not found at: {}'.format(filename)
        with open(filename, 'rb') as f:
            box_list        = pickle.load(f)
        return self.create_roidb_from_box_list(box_list, gt_roidb)

    def _get_comp_id(self):
        comp_id             = (self._comp_id + '_' + self._salt if self.config['use_salt'] else self._comp_id)
        return comp_id
    def _get_kitti_results_file_template(self):
        """
        Use this function to generate filename template for each result file
        """
        result_dir          = self.result_path
        if not osp.exists(result_dir):
            os.makedirs(result_dir)
        path                = osp.join(result_dir, "{:s}.txt")
        return path, result_dir
    def _write_kitti_results_file(self, all_boxes):
        # write the kitti detection results into file
        
        filename_tem, _     = self._get_kitti_results_file_template()
        for im_ind, index in enumerate(self._image_index):
            # one by one write the result
            filename        = filename_tem.format(index)
            with open(filename, 'w') as f:
                for cls_ind, cls in enumerate(self.classes):
                    if cls == '__background__':
                        continue
                    dets    = all_boxes[cls_ind][im_ind]
                    if dets == []:
                        continue
                    # the kitti expects 0-based indices
                    for k in range(dets.shape[0]):
                        f.write('{:s} {:d} {:d} {:.3f} {:.1f} {:.1f} {:.1f} {:.1f}\
                                {:d} {:d} {:d} {:d} {:d} {:d} {:.3f} {:.3f}\n'.format(cls, -1, -1, 0.0, dets[k,0], dets[k,1], dets[k,2], dets[k,3], -1,-1,-1,-1,-1,-1,0.0,dets[k,-1]))

    def _do_cpp_eval(self, gtFileDir=None, detFileDir=None, nameListFilePath=None):
        """
        Use the modification of kitti evaluation cpp to eval the detections
        """
        if gtFileDir is None:
            gtFileDir       = osp.join(self._data_path, 'label_2')
        if detFileDir is None:
            _, detFileDir   = self._get_kitti_results_file_template()
        if nameListFilePath is None:
            nameListFilePath    = self._image_index_file 
        print("Using kitti official evaluation benchmark")
        command         = '/home/tzz/projects/dataset/kitti/detection/devkit/cpp/evaluate_object {} {} {}'.format(nameListFilePath, gtFileDir, detFileDir)
        print("Excuting command: {}".format(command))
        os.system(command)
        print("Excuted end")

    def evaluate_detections(self, all_boxes):
        """
        wrapper for using evaluations 
        """
        # must do self._write_kitti_results_file firstly
        self._write_kitti_results_file(all_boxes)
        self._do_cpp_eval()

    def competition_mode(self, on):
        if on:
            self.config['use_salt'] = False
            self.config['cleanup']  = False
        else:
            self.config['use_salt'] = True
            self.config['cleanup']  = True

